package com.zyf.project.kvdb.client.core.impl;

import com.zyf.project.binaryproto.common.BinaryProtocol;
import com.zyf.project.kvdb.client.core.KVDBOperator;
import com.zyf.project.kvdb.client.partitioner.Partitioner;
import com.zyf.project.kvdb.client.partitioner.impl.SimpleMurmur3HashPartitioner;
import com.zyf.project.kvdb.protocol.exception.KVDBException;
import com.zyf.project.kvdb.protocol.handler.impl.NettyClient;
import com.zyf.project.kvdb.protocol.proto.wal.Entity;
import com.zyf.project.kvdb.protocol.proto.wal.KV;
import com.zyf.project.kvdb.protocol.proto.wal.impl.KVItem;
import com.zyf.project.kvdb.protocol.proto.wal.impl.WalEntity;
import com.zyf.project.kvdb.wal.config.WalConfig;
import com.zyf.project.kvdb.wal.logger.FileLogger;
import com.zyf.project.kvdb.wal.walinf.Wal;
import com.zyf.project.kvdb.wal.walinf.WalDataCoder;
import com.zyf.project.kvdb.wal.walinf.WalIterator;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import utils.binary.impl.Bytes;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.*;

/**
 * @author Zhou Yifan
 */
public class KVDBCluster implements KVDBOperator {

    private static final Logger LOGGER = LoggerFactory.getLogger(KVDBCluster.class);

    /**
     * 线程池，用于异步多数据的数据操作
     */
    private ExecutorService executorService;

    /**
     * 服务器分片器
     */
    private Partitioner partitioner;

    /**
     * 数据库名称
     */
    private String dbName;

    /**
     * 集群节点的数据库实例的实际操作对象
     */
    private KVDBOperator[] operators;

    /**
     * 是否是批处理模式
     */
    private Boolean batchMode = false;

    /**
     * 批处理数据
     */
    private Map<Bytes, Bytes> batch = new ConcurrentHashMap<>();

    /**
     * WAL
     */
    private final Wal<Entity> wal;

    /**
     * 是否需要 WAL redo 操作
     */
    private volatile boolean needRedo = true;

    @Override
    public String toString() {
        return "KVDBCluster{" +
                "dbName='" + dbName + '\'' +
                ", clusterNum=" + operators.length +
                ", batchMode=" + batchMode +
                ", needRedo=" + needRedo +
                '}';
    }

    public Map<Bytes, Bytes> getBatch() {
        return batch;
    }

    public KVDBCluster(String db, KVDBOperator[] operators) throws KVDBException, IOException {
        if (null != operators && operators.length > 0) {
            this.dbName = db;
            this.operators = operators;
            this.wal = new FileLogger<>(new WalConfig(this.dbName, 1, true), new WalDataCoder<Entity>() {
                @Override
                public byte[] encode(Entity data) {
                    return BinaryProtocol.encode(data, Entity.class);
                }

                @Override
                public Entity decode(byte[] data) {
                    return BinaryProtocol.decode(data, Entity.class);
                }
            });
            init();
        } else {
            throw new KVDBException("no cluster config present");
        }
    }

    public KVDBCluster(String db, NettyClient[] clients) throws KVDBException, IOException {
        if (null != clients && clients.length > 0) {
            this.dbName = db;
            this.operators = new KVDBSingle[clients.length];
            int i = 0;
            for (NettyClient client : clients) {
                operators[i] = new KVDBSingle(client);
                i++;
            }
            this.wal = new FileLogger<>(new WalConfig(this.dbName, 1, true), new WalDataCoder<Entity>() {
                @Override
                public byte[] encode(Entity data) {
                    return BinaryProtocol.encode(data, Entity.class);
                }

                @Override
                public Entity decode(byte[] data) {
                    return BinaryProtocol.decode(data, Entity.class);
                }
            });
            init();
        } else {
            throw new KVDBException("no cluster config present");
        }
    }

    /**
     * 一个简单的初始化函数，目的就是简化代码
     *
     * @throws KVDBException KVDB 可能的异常
     */
    private void init() throws KVDBException {
        partitioner = new SimpleMurmur3HashPartitioner(operators.length);
        executorService = Executors.newFixedThreadPool(operators.length);
        checkAndRedoWal();
    }

    /**
     * 检查是否需要 redo，并执行相关操作
     *
     * @throws KVDBException KVDB 可能的错误
     */
    private void checkAndRedoWal() throws KVDBException {
        if (!needRedo) {
            return;
        }
        synchronized (wal) {
            try {
                WalIterator<Entity> iterator = wal.forwardIterator();
                while (iterator.hasNext()) {
                    Entity entity = iterator.next();
                    for (KV kv : entity.getKVs()) {
                        if (!operators[partitioner.partition(kv.getKey())].put(new Bytes(kv.getKey()), new Bytes(kv.getValue()), false)) {
                            throw new KVDBException("redo failed");
                        }
                    }
                }
                wal.checkpoint();
                needRedo = false;
            } catch (IOException e) {
                LOGGER.error("redo error in " + dbName, e);
                throw new KVDBException(e);
            }
        }
    }

    /**
     * 存在性查询
     *
     * @param key 待查询的键
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean exists(Bytes key) throws KVDBException {
        checkAndRedoWal();
        return operators[partitioner.partition(key)].exists(key);
    }

    /**
     * 存在性查询，支持多个键
     *
     * @param keys 待查询的键
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean[] exists(Bytes @NotNull ... keys) throws KVDBException {
        checkAndRedoWal();
        try {
            CompletionService<ExecuteResultInBatch<Boolean>> completionService = new ExecutorCompletionService<>(executorService);
            for (int i = 0; i < keys.length; i++) {
                final int index = i;
                completionService.submit(() -> {
                    try {
                        return new ExecuteResultInBatch<>(index, operators[partitioner.partition(keys[index])].exists(keys[index]));
                    } catch (KVDBException e) {
                        return new ExecuteResultInBatch<>(index, e);
                    }
                });
            }
            boolean[] results = new boolean[keys.length];
            for (int i = 0; i < keys.length; i++) {
                ExecuteResultInBatch<Boolean> result = completionService.take().get();
                if (result.failed()) {
                    throw result.exception;
                }
                results[result.index] = result.data;
            }
            return results;
        } catch (ExecutionException | InterruptedException e) {
            LOGGER.error("exists error", e);
            Thread.currentThread().interrupt();
            throw new KVDBException(e.getMessage());
        }
    }

    /**
     * 键值对获取
     *
     * @param key 键，{@link Bytes} 类型
     * @return 键对应的值，{@link Bytes} 类型
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public Bytes get(Bytes key) throws KVDBException {
        checkAndRedoWal();
        return operators[partitioner.partition(key)].get(key);
    }

    /**
     * 键值对获取
     *
     * @param keys 键，{@link Bytes} 类型数组
     * @return 键对应的值，{@link Bytes} 类型数组
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public Bytes[] get(Bytes @NotNull ... keys) throws KVDBException {
        checkAndRedoWal();
        try {
            CompletionService<ExecuteResultInBatch<Bytes>> completionService = new ExecutorCompletionService<>(executorService);
            for (int i = 0; i < keys.length; i++) {
                final int index = i;
                completionService.submit(() -> {
                    try {
                        return new ExecuteResultInBatch<>(index, operators[partitioner.partition(keys[index])].get(keys[index]));
                    } catch (KVDBException e) {
                        return new ExecuteResultInBatch<>(index, e);
                    }
                });
            }
            Bytes[] results = new Bytes[keys.length];
            for (int i = 0; i < keys.length; i++) {
                ExecuteResultInBatch<Bytes> result = completionService.take().get();
                if (result.failed()) {
                    throw result.exception;
                }
                results[result.index] = result.data;
            }
            return results;
        } catch (ExecutionException | InterruptedException e) {
            LOGGER.error("get error", e);
            Thread.currentThread().interrupt();
            throw new KVDBException(e.getMessage());
        }
    }

    /**
     * 设置键值对（添加）
     *
     * @param key   键
     * @param value 值
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean put(Bytes key, Bytes value) throws KVDBException {
        return put(key, value, false);
    }

    /**
     * 设置键值对（添加）
     *
     * @param key   键
     * @param value 值
     * @param async 是否异步（Promise 消息）
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean put(Bytes key, Bytes value, boolean async) throws KVDBException {
        checkAndRedoWal();
        synchronized (batchMode) {
            boolean result = operators[partitioner.partition(key)].put(key, value, async);
            if (Boolean.TRUE.equals(batchMode)) {
                batch.put(key, value);
            }
            return false;
        }
    }

    /**
     * 开启批处理
     *
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean batchBegin() throws KVDBException {
        checkAndRedoWal();
        try {
            synchronized (batchMode) {
                batch.clear();
                CompletionService<ExecuteResultInBatch<Boolean>> completionService = new ExecutorCompletionService<>(executorService);
                for (int i = 0; i < operators.length; i++) {
                    final int index = i;
                    completionService.submit(() -> {
                        try {
                            return new ExecuteResultInBatch<>(index, operators[index].batchBegin());
                        } catch (KVDBException e) {
                            return new ExecuteResultInBatch<>(index, e);
                        }
                    });
                }
                for (int i = 0; i < operators.length; i++) {
                    ExecuteResultInBatch<Boolean> result = completionService.take().get();
                    if (result.failed()) {
                        throw result.exception;
                    }
                    if (Boolean.FALSE.equals(result.data)) {
                        return false;
                    }
                }
                batchMode = true;
                return true;
            }
        } catch (ExecutionException | InterruptedException e) {
            LOGGER.error("batch begin error!", e);
            Thread.currentThread().interrupt();
            throw new KVDBException(e.getMessage());
        }
    }

    /**
     * 取消批处理
     *
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean batchAbort() throws KVDBException {
        checkAndRedoWal();
        try {
            synchronized (batchMode) {
                batch.clear();
                CompletionService<ExecuteResultInBatch<Boolean>> completionService = new ExecutorCompletionService<>(executorService);
                for (int i = 0; i < operators.length; i++) {
                    final int index = i;
                    completionService.submit(() -> {
                        try {
                            return new ExecuteResultInBatch<>(index, operators[index].batchAbort());
                        } catch (KVDBException e) {
                            return new ExecuteResultInBatch<>(index, e);
                        }
                    });
                }
                for (int i = 0; i < operators.length; i++) {
                    ExecuteResultInBatch<Boolean> result = completionService.take().get();
                    if (result.failed()) {
                        throw result.exception;
                    }
                    if (Boolean.FALSE.equals(result.data)) {
                        return false;
                    }
                }
                batchMode = false;

                return true;
            }
        } catch (ExecutionException | InterruptedException e) {
            LOGGER.error("batch begin error!", e);
            Thread.currentThread().interrupt();
            throw new KVDBException(e.getMessage());
        }
    }

    /**
     * 提交批处理，服务器掉线重连后会丢失未提交的数据，未提交的 batch 对其他客户端连接而言是不可见的
     *
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean batchCommit() throws KVDBException {
        checkAndRedoWal();
        try {
            synchronized (batchMode) {
                if (Boolean.FALSE.equals(batchMode)) {
                    return false;
                }
                batchMode = false;
                if (batch.size() == 0) {
                    throw new KVDBException("error batch size");
                }
                synchronized (wal) {
                    KVItem[] walkvs = new KVItem[batch.size()];
                    int[] ps = new int[partitioner.getPartitionCount()];
                    int j = 0;
                    for (Map.Entry<Bytes, Bytes> entry : batch.entrySet()) {
                        ps[partitioner.partition(entry.getKey())]++;
                        walkvs[j] = new KVItem(entry.getKey().toBytes(), entry.getValue().toBytes());
                        j++;
                    }
                    batch.clear();
                    // append wal
                    try {
                        wal.append(WalEntity.newPutEntity(walkvs));
                    } catch (IOException e) {
                        LOGGER.error("wal append error", e);
                        throw new KVDBException(e.getMessage());
                    }
                    // sub batch commit, any error need redo
                    try {
                        CompletionService<ExecuteResultInBatch<Boolean>> completionService = new ExecutorCompletionService<>(executorService);
                        for (int i = 0; i < operators.length; i++) {
                            final int index = i;
                            completionService.submit(() -> {
                                try {
                                    if (ps[index] > 0) {
                                        return new ExecuteResultInBatch<>(index, operators[index].batchCommit());
                                    } else {
                                        return new ExecuteResultInBatch<>(index, true);
                                    }
                                } catch (KVDBException e) {
                                    return new ExecuteResultInBatch<>(index, e);
                                }
                            });
                        }
                        for (int i = 0; i < operators.length; i++) {
                            ExecuteResultInBatch<Boolean> result = completionService.take().get();
                            if (result.failed()) {
                                throw result.exception;
                            }
                            if (Boolean.FALSE.equals(result.data)) {
                                throw new KVDBException("batch commit error");
                            }
                        }
                    } catch (ExecutionException | InterruptedException e) {
                        needRedo = true;
                        LOGGER.error("batch commit error!", e);
                        Thread.currentThread().interrupt();
                        throw new KVDBException(e.getMessage());
                    }
                    // wal checkpoint
                    try {
                        wal.checkpoint();
                    } catch (IOException e) {
                        LOGGER.error("wal checkpoint error", e);
                    }

                    return true;
                }
            }
        } finally {
            batchAbort();
        }
    }

    /**
     * 提交批处理
     *
     * @param size 此次批处理操作去重后的 key 数量
     * @return 布尔值
     * @throws KVDBException KVDB 可能存在的异常
     */
    @Override
    public boolean batchCommit(long size) throws KVDBException {
        checkAndRedoWal();
        try {
            synchronized (batchMode) {
                batchMode = false;
                synchronized (wal) {
                    if (size <= 0 || batch.size() == 0) {
                        throw new KVDBException("error batch size");
                    }
                    if (size != batch.size()) {
                        throw new KVDBException("batch commit size wrong, expected: " + size + ", actually: " + batch.size());
                    }
                    KVItem[] walkvs = new KVItem[batch.size()];
                    int[] ps = new int[partitioner.getPartitionCount()];
                    int j = 0;
                    for (Map.Entry<Bytes, Bytes> entry : batch.entrySet()) {
                        ps[partitioner.partition(entry.getKey())]++;
                        walkvs[j] = new KVItem(entry.getKey().toBytes(), entry.getValue().toBytes());
                        j++;
                    }
                    batch.clear();
                    // wal append
                    try {
                        wal.append(WalEntity.newPutEntity(walkvs));
                    } catch (IOException e) {
                        LOGGER.error("wal append error", e);
                        throw new KVDBException(e.getMessage());
                    }
                    // sub batch commit, any error need redo
                    try {
                        CompletionService<ExecuteResultInBatch<Boolean>> completionService = new ExecutorCompletionService<>(executorService);
                        for (int i = 0; i < operators.length; i++) {
                            final int index = i;
                            completionService.submit(() -> {
                                try {
                                    if (ps[index] > 0) {
                                        return new ExecuteResultInBatch<>(index, operators[index].batchCommit(ps[index]));
                                    } else {
                                        return new ExecuteResultInBatch<>(index, true);
                                    }
                                } catch (KVDBException e) {
                                    return new ExecuteResultInBatch<>(index, e);
                                }
                            });
                        }
                        for (int i = 0; i < operators.length; i++) {
                            ExecuteResultInBatch<Boolean> result = completionService.take().get();
                            if (result.failed()) {
                                throw result.exception;
                            }
                            if (Boolean.FALSE.equals(result.data)) {
                                throw new KVDBException("batch commit error");
                            }
                        }
                    } catch (ExecutionException | InterruptedException e) {
                        needRedo = true;
                        LOGGER.error("batch commit error!", e);
                        Thread.currentThread().interrupt();
                        throw new KVDBException(e.getMessage());
                    }

                    // wal checkpoint
                    try {
                        wal.checkpoint();
                    } catch (IOException e) {
                        LOGGER.error("wal checkpoint error", e);
                    }

                    return true;
                }
            }
        } finally {
            batchAbort();
        }
    }

    /**
     * 关闭数据库
     */
    @Override
    public void close() {
        for (KVDBOperator operator : operators) {
            operator.close();
        }
        if (null != executorService && !executorService.isShutdown()) {
            executorService.shutdown();
        }
        if (null != wal) {
            try {
                wal.close();
            } catch (Exception e) {
                LOGGER.error("wal close error", e);
            }
        }
    }

    /**
     * <h1>批处理执行结果</h1>
     *
     * @param <E> 批处理结果类型
     */
    static class ExecuteResultInBatch<E> {

        /**
         * 序号
         */
        private int index;
        /**
         * 处理结果数据
         */
        private E data;
        /**
         * 处理结果异常
         */
        private KVDBException exception;

        private ExecuteResultInBatch() {
        }

        private ExecuteResultInBatch(int index, E data) {
            this.index = index;
            this.data = data;
        }

        private ExecuteResultInBatch(int index, KVDBException e) {
            this.index = index;
            this.exception = e;
        }

        /**
         * 判断是否处理成功，即没有异常就是成功
         *
         * @return 成功则为 true
         */
        public boolean failed() {
            return null != exception;
        }
    }
}
